import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 5)))
sys.path.insert(0, parent_path)

# from queue import Queue
from core.task.modules.processors.base_processor import BaseProcessor
from typing import *
import numpy as np
import cv2, time, traceback
from threading import Thread
from core.utils.logger import logging

import subprocess


class SteamPush(BaseProcessor):
    """任务管道的拉流组件
    跟随任务启动后, 会在task_init时指定拉流地址"""
    def __init__(self, keys, 
                 SteamPush_url = None, 
                 SteamPush_size = None,
                 SteamPush_fps = 5,
                 ) -> None:
        """
        input:
            keys:{
                in: 需要推流数据的键值
            }
        """
        super().__init__(keys)
        self.url = SteamPush_url
        self.width = SteamPush_size[0] if not SteamPush_size is None else None
        self.height = SteamPush_size[1] if not SteamPush_size is None else None
        self.push_image = np.zeros(shape=[self.height,self.width,3],dtype=np.uint8) if not SteamPush_size is None else None
        self.fps = SteamPush_fps
        self.wait_time = 1/self.fps
        self.dymamic_frame_idx = 0
        self.command = None
        self.pipe = None
        self.run_thread = None
    
    def _init_check(self):
        assert not self.url is None
        assert not self.width is None
        assert not self.height is None
        assert not self.fps is None
        assert not self.command is None
        assert not self.pipe is None

    def _set(self, data: Dict):
        if data.get("SteamPush_url",False):
            self.url = data["SteamPush_url"]
        if data.get("SteamPush_size",False):
            self.width = data["SteamPush_size"][0]
            self.height = data["SteamPush_size"][1]
        if data.get("SteamPush_fps",False):
            self.fps = data["SteamPush_fps"]
            self.wait_time = 1/self.fps
        self.push_image = np.zeros(shape=[self.height,self.width,3],dtype=np.uint8)
        self.command = self._build_command()
        self.pipe = subprocess.Popen(self.command,stdin=subprocess.PIPE,)
        self.run_thread = Thread(target=self.run_loop)

    
    def run_loop(self):
        try:
            while True:
                time.sleep(1)
                if self.is_running():
                    break
                elif self.is_close():
                    break
                logging.info("{} waitting start....".format(self.__class__.__name__))
                
            while self.status == 2:
                ret = self.put(self.push_image)
                if not ret:
                    raise Exception()
                # time.sleep(max(self.wait_time - (time.time() - start_time), 0))
                self.dymamic_frame_idx+=1
                time.sleep(self.wait_time)
        except Exception:
            logging.error(traceback.format_exc())
            self.close()
        finally:
            logging.info("|SteamPush.video_read_loop done|")

    def dymamic_fps_loop(self):
        try:
            while True:
                time.sleep(1)
                if self.is_running():
                    break
                elif self.is_close():
                    break

            start_time = time.time()
            self.dymamic_frame_idx = 0
            while self.status == 2:
                if self.dymamic_frame_idx>=self.fps*10:
                    fps = self.dymamic_frame_idx/(time.time()-start_time)
                    if fps/self.fps<0.95:
                        self.wait_time*=max(fps/self.fps,0.8)
                    elif fps/self.fps>1.01:
                        self.wait_time*=min(fps/self.fps,1.2)
                    else:
                        pass
                    # print("now fps: {}, scale: {}\n".format(fps,fps/self.fps))
                    self.dymamic_frame_idx = 0
                    start_time = time.time()
                time.sleep(5)
        except Exception:
            logging.error(traceback.format_exc())
            self.close()

    def put(self,data):
        try:
            self.pipe.stdin.write(data.tostring())
            return True
        except Exception as e:
            logging.error("|RtmpStream put fail|")
            logging.error(e)
            return False

    def _build_command(self):
        return ['ffmpeg',
                '-y','-an',
                '-f', 'rawvideo',
                '-vcodec', 'rawvideo',
                '-pix_fmt', 'bgr24',
                '-s', "{}x{}".format(self.width, self.height),
                '-r', str(self.fps),
                '-i', '-',
                '-c:v', 'libx264',
                '-b:v', "2000k",
                '-pix_fmt', 'yuv420p',
                '-preset', 'fast',
                '-tune', 'zerolatency',
                '-f', 'flv',
                self.url]
    
    def _call(self, data: Dict) -> Dict:
        self.push_image = cv2.resize(data[self.keys["in"]],[self.width,self.height])
        return data

    def _start(self):
        self.run_thread.start()
        Thread(target=self.dymamic_fps_loop).start()


    def close(self):
        try:
            if self.status == 2:
                super().close()
                self.pipe.terminate()
                self.pipe.communicate()
            else:
                logging.warning("|SteamPush status is closed|")
        except Exception as e:
            raise e

    
if __name__ == "__main__":
    from core.task.modules.processors import SteamPull
    steam_puller = SteamPull(
        keys={
            "out":"pull_data"
            },
        SteamPull_url = "rtmp://218.94.155.86:1936/live/livestream_74_1",
        SteamPull_size=[1280,720],
    )
    steam_puller.set()
    steam_puller.start()
    steam_pusher = SteamPush(
        keys={
            "in":"pull_data"
            },
        SteamPush_url = "rtmp://218.94.155.86:1936/live/livestream_74_2",
        SteamPush_size=[1280,720],
    )
    steam_pusher.set()
    steam_pusher.start()
    start_t = time.time()
    while True:
        # if time.time() - start_t>5:
        #     break
        databox = {}
        databox = steam_puller(databox)
        databox["pull_data"] = databox["pull_data"][0]
        data = steam_pusher(databox)
        cv2.imwrite("temp.jpg",data["pull_data"])
        time.sleep(0.1)

    steam_pusher.close()
    print("push close")

    while True:
        time.sleep(1)
